// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
//
// Miscellaneous number encoding/decoding routines
// - Varint coding
// - ZigZag coding

#ifndef FACEBOOK_COMMON_ENCODE_CODING_H_
#define FACEBOOK_COMMON_ENCODE_CODING_H_

#include <melon/group_varint.h>
#include <melon/likely.h>
#include <algorithm>
#include <utility>
#include <pollux/common/encode/uInt128.h>
#include <pollux/common/strings/byte_stream.h>

namespace kumo {
    // Variable-length integer encoding, using a little-endian, base-128
    // representation.
    // The MSb is set on all bytes except the last.
    class Varint {
    public:
        // Max size for a k-bit value: 1 + floor(k/7) bytes
        enum {
            kMaxSize32 = 5,
            kMaxSize64 = 10,
            kMaxSize128 = 19,
        };

        // Make <a,b> contain (a<<bBits) + b, if b has less than n bits
        // (that is, b < 2^n)
        // This is a useful transformation for encoding a 64-bit value together
        // with a small value using fewer bytes than encoding a and b separately.
        static UInt128 shift(uint64_t a, uint64_t b, uint32_t bBits) {
            return (UInt128(a) << bBits) | b;
        }

        static std::pair<uint64_t, uint64_t> unshift(UInt128 val, uint32_t bBits) {
            std::pair<uint64_t, uint64_t> p;
            p.second = val.lo() & ((static_cast<uint64_t>(1) << bBits) - 1);
            val >>= bBits;
            p.first = val.lo();
            return p;
        }

        // Determine the number of bytes needed to represent "val".
        // 32-bit values need at most 5 bytes.
        // 64-bit values need at most 10 bytes.
        static int32_t size(uint64_t val) {
            int32_t s = 1;
            while (val >= 128) {
                ++s;
                val >>= 7;
            }
            return s;
        }

        // Determine the number of bytes needed to represent "val".
        // 32-bit values need at most 5 bytes.
        // 64-bit values need at most 10 bytes.
        // 128-bit values need at most 19 bytes.
        static int32_t size128(UInt128 val) {
            int32_t s = 1;
            while (val.hi() != 0 || val.lo() >= 128) {
                ++s;
                val >>= 7;
            }
            return s;
        }

        // Encode "val" into the buffer pointed to by *dest, and advance *dest.
        // No bounds checking is done.
        static void encode(uint64_t val, char **dest) {
            char *p = *dest;
            while (val >= 128) {
                *p++ = 0x80 | (static_cast<char>(val) & 0x7f);
                val >>= 7;
            }
            *p++ = static_cast<char>(val);
            *dest = p;
        }

        static void encode128(UInt128 val, char **dest) {
            char *p = *dest;
            while (val.hi() != 0 || val.lo() >= 128) {
                *p++ = 0x80 | (static_cast<char>(val.lo()) & 0x7f);
                val >>= 7;
            }
            *p++ = static_cast<char>(val.lo());
            *dest = p;
        }

        // Encode "val" into the byte sink pointed-to by "sink".
        static void encodeToByteSink(uint64_t val, strings::ByteSink *sink) {
            char buf[kMaxSize64];
            char *p = buf;
            encode(val, &p);
            sink->append(melon::StringPiece(buf, p - buf));
        }

        static void encode128ToByteSink(UInt128 val, strings::ByteSink *sink) {
            char buf[kMaxSize128];
            char *p = buf;
            encode128(val, &p);
            sink->append(melon::StringPiece(buf, p - buf));
        }

        // Returns true if decode can be called without causing a KCHECK failure.
        // The pointers are not adjusted at all
        static bool canDecode(melon::StringPiece src) {
            src = src.subpiece(0, kMaxSize64);
            return std::any_of(
                src.begin(), src.end(), [](char v) { return ~v & 0x80; });
        }

        // Decode a value from the buffer pointed to by *src, and advance *src.
        // Assumes that *src has at least max_size bytes available.
        static uint64_t decode(const char **src, int64_t max_size) {
            if (max_size >= kMaxSize64) {
                // Fast Path
                uint64_t x;
                int64_t b;
                const signed char *p = reinterpret_cast<const signed char *>(*src);
                do {
        // clang-format off
        b = *p++; x  = (b & 0x7f)      ; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) <<  7; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 14; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 21; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 28; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 35; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 42; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 49; if (b >= 0) break;
        b = *p++; x |= (b & 0x7f) << 56; if (b >= 0) break;
        b = *p++; x |= static_cast<uint64_t>(b & 0x7f) << 63; if (b >= 0) break;
                    // clang-format on
                    KCHECK(false); // varints > 64bits == data corruption
                } while (false);
                *src = reinterpret_cast<const char *>(p);
                return x;
            }
            const char *p = *src;
            uint64_t val = 0;
            int32_t shift = 0;
            while (*p & 0x80) {
                KCHECK_GT(max_size, 1); // We must have room for the last byte, too
                --max_size;
                val |= static_cast<uint64_t>(*p++ & 0x7f) << shift;
                shift += 7;
            }
            val |= static_cast<uint64_t>(*p++) << shift;
            *src = p;
            return val;
        }

        static UInt128 decode128(const char **src, int64_t max_size) {
            if (max_size > kMaxSize128) {
                // Varint-encoded numbers are at most 19 bytes, and we want to catch
                // the case where we have more than 18 consecutive bytes with bit 7 set
                // (which would be an error)
                max_size = kMaxSize128;
            }
            const char *p = *src;
            UInt128 val = 0;
            int32_t shift = 0;
            while (*p & 0x80) {
                KCHECK_GT(max_size, 1); // We must have room for the last byte, too
                --max_size;
                val |= UInt128(*p++ & 0x7f) << shift;
                shift += 7;
            }
            val |= UInt128(*p++) << shift;
            *src = p;
            return val;
        }

        // Decode a value from a StringPiece, and advance the StringPiece.
        static uint64_t decode(melon::StringPiece *data) {
            const char *p = data->start();
            uint64_t val = decode(&p, data->size());
            data->advance(p - data->start());
            return val;
        }

        static UInt128 decode128(melon::StringPiece *data) {
            const char *p = data->start();
            UInt128 val = decode128(&p, data->size());
            data->advance(p - data->start());
            return val;
        }

        // Decode a value from a ByteSource, and advance the ByteSource.
        static uint64_t decodeFromByteSource(strings::ByteSource *src) {
            uint64_t val = 0;
            int32_t shift = 0;
            int32_t max_size = kMaxSize64;
            melon::StringPiece chunk;
            int32_t remaining = 0;
            const char *p = nullptr;
            for (;;) {
                if (remaining == 0) {
                    KCHECK(src->next(&chunk));
                    p = chunk.start();
                    remaining = chunk.size();
                    DKCHECK_GT(remaining, 0);
                }
                --remaining;
                if (*p & 0x80) {
                    KCHECK_GT(max_size, 1); // We must have room for the last byte, too
                    --max_size;
                    val |= static_cast<uint64_t>(*p++ & 0x7f) << shift;
                    shift += 7;
                } else {
                    val |= static_cast<uint64_t>(*p++) << shift;
                    break;
                }
            }
            if (remaining) {
                src->backUp(remaining);
            }
            return val;
        }

        static UInt128 decode128FromByteSource(strings::ByteSource *src) {
            UInt128 val = 0;
            int32_t shift = 0;
            int32_t max_size = kMaxSize128;
            melon::StringPiece chunk;
            int32_t remaining = 0;
            const char *p = nullptr;
            for (;;) {
                if (remaining == 0) {
                    KCHECK(src->next(&chunk));
                    p = chunk.start();
                    remaining = chunk.size();
                    DKCHECK_GT(remaining, 0);
                }
                --remaining;
                if (*p & 0x80) {
                    KCHECK_GT(max_size, 1); // We must have room for the last byte, too
                    --max_size;
                    val |= UInt128(*p++ & 0x7f) << shift;
                    shift += 7;
                } else {
                    val |= UInt128(*p++) << shift;
                    break;
                }
            }
            if (remaining) {
                src->backUp(remaining);
            }
            return val;
        }
    };

    // Zig-zag encoding that maps signed integers with a small absolute value
    // to unsigned integers with a small (positive) value.
    // if x >= 0, ZigZag::encode(x) == 2*x
    // if x < 0, ZigZag::encode(x) == -2*x + 1
    class ZigZag {
    public:
        static uint64_t encode(int64_t val) {
            // Bit-twiddling magic stolen from the Google protocol buffer document;
            // val >> 63 is an arithmetic shift because val is signed
            return (static_cast<uint64_t>(val) << 1) ^ (val >> 63);
        }

        static __uint128_t encodeInt128(__int128_t val) {
            return (static_cast<__uint128_t>(val) << 1) ^ (val >> 127);
        }

        template<typename U, typename T = typename std::make_signed<U>::type>
        static T decode(U val) {
            return static_cast<T>((val >> 1) ^ -(val & 1));
        }
    };

    namespace internal {
        class ByteSinkAppender {
        public:
            /* implicit */
            ByteSinkAppender(strings::ByteSink *out) : out_(out) {
            }

            void operator()(melon::StringPiece sp) {
                out_->append(sp.data(), sp.size());
            }

        private:
            strings::ByteSink *out_;
        };
    } // namespace internal

    // Import GroupVarint encoding / decoding code from melon

    typedef melon::GroupVarint32 GroupVarint32;
    typedef melon::GroupVarint64 GroupVarint64;

    typedef melon::GroupVarintEncoder<uint32_t, internal::ByteSinkAppender>
    GroupVarint32Encoder;
    typedef melon::GroupVarintEncoder<uint64_t, internal::ByteSinkAppender>
    GroupVarint64Encoder;

    typedef melon::GroupVarintDecoder<uint32_t> GroupVarint32Decoder;
    typedef melon::GroupVarintDecoder<uint64_t> GroupVarint64Decoder;
} // namespace kumo

#endif /* FACEBOOK_COMMON_ENCODE_CODING_H_ */
